from prediction import *


class Strategy:
    # 修改Strategy类的初始化方法和选股方法，使其使用可用股票列表

    def __init__(self, data_days=10):
        """
        Strategy
        @data_days: 回溯交易日天数
        """
        self.base_data_path = './data/'
        self.data_path = './data/stocks/'
        self.model_path = './models/'
        self.data_days = data_days
        
        # 读取可用股票列表而不是原始股票列表
        available_stocks_file = '{}hs300_stocks_available.csv'.format(self.base_data_path)
        if os.path.exists(available_stocks_file):
            self.stocks = pd.read_csv(available_stocks_file)
        else:
            # 如果可用股票列表不存在，创建它
            stock_files = []
            if os.path.exists(self.data_path):
                stock_files = [f.replace('.csv', '') for f in os.listdir(self.data_path) if f.endswith('.csv')]
            
            available_stocks = [{'code': code} for code in stock_files if code != 'sh.000300']
            self.stocks = pd.DataFrame(available_stocks)
            self.stocks.to_csv(available_stocks_file, index=False)
        
        self.stocks_codes = self.stocks['code']

    # 然后需要修改所有的选股方法，示例修改choose_by_bm:
    def choose_by_bm(self, today: tuple, num=10):
        """选择账面市值比(BM)最高的num只股票"""
        # 获取所有股票当日情况
        codes = self.stocks_codes
        if len(codes) < num:
            num = len(codes)  # 确保不超过可用股票数量
            
        # 计算所有可用股票的BM因子
        bm_factors = {}
        for stock_code in codes:
            try:
                file_path = '{}{}.csv'.format(self.data_path, stock_code)
                if not os.path.exists(file_path):
                    continue
                    
                stock_data = pd.read_csv(file_path, index_col='date')
                if today[0] in stock_data.index:
                    # 计算BM因子
                    bm = stock_data.loc[today[0], 'pb']
                    if not pd.isna(bm) and bm != 0:
                        bm_factors[stock_code] = 1 / bm
            except Exception as e:
                print(f"计算股票 {stock_code} BM因子时出错: {str(e)}")
                continue
        
        # 排序并选择因子最高的num只股票
        sorted_stocks = sorted(bm_factors.items(), key=lambda x: x[1], reverse=True)
        chosen_stocks = pd.Series([item[0] for item in sorted_stocks[:num]])
        return chosen_stocks

    def choose_by_mf(self, today: tuple, number: int):
        """
        选择最近data_days中动量因子(Momentum Factor)最高的number只股票
        """
        # 第一次买入策略应大于策略所需数据天数
        if today[1] < self.data_days:
            return pd.Series(None)
        # 建立用于计算平均MF的DF
        stocks_data = pd.DataFrame(self.stocks_codes)
        stocks_data['aver_MF'] = 0
        stocks_data = stocks_data.set_index('code')
        # 到交易日前一日为止共data_days日期序号
        days = range(today[1] - self.data_days, today[1])
        for stock_code in self.stocks_codes:
            sum_MF = 0
            valid_days = self.data_days
            stock_data = pd.read_csv('{}{}.csv'.format(self.data_path, stock_code), index_col='date')
            for day in days:
                if self.trading_dates[day] in stock_data.index:
                    # 加入收益率
                    pc = stock_data.loc[self.trading_dates[day], 'preclose']
                    close = stock_data.loc[self.trading_dates[day], 'close']
                    sum_MF += close / pc
                else:
                    valid_days -= 1
            if valid_days != 0:
                aver_MF = sum_MF / valid_days
            else:
                aver_MF = 0
            if aver_MF > 0:
                stocks_data.loc[stock_code, 'aver_MF'] = aver_MF
        # print(stocks_data)
        stocks_data.sort_values(by='aver_MF', ascending=False, inplace=True)
        # print(stocks_data.index)
        if len(stocks_data.index) > number:
            # 取0到number-1共number只股票
            return stocks_data.index[0:number]
        else:
            # 取全部股票
            return stocks_data.index[:]

    def choose_by_tr(self, today: tuple, number: int):
        """
        选择最近data_days中换手率因子(Unusual Turnover Rate, 异常换手率)最高的number只股票
        """
        # 第一次买入策略应大于策略所需数据天数
        if today[1] < self.data_days:
            return pd.Series(None)
        # 建立用于计算平均TR的DF
        stocks_data = pd.DataFrame(self.stocks_codes)
        stocks_data['aver_TR'] = 0
        stocks_data = stocks_data.set_index('code')
        # 到交易日前一日为止共data_days日期序号
        days = range(today[1] - self.data_days, today[1])
        for stock_code in self.stocks_codes:
            sum_TR = 0
            valid_days = self.data_days
            stock_data = pd.read_csv('{}{}.csv'.format(self.data_path, stock_code), index_col='date')
            for day in days:
                if self.trading_dates[day] in stock_data.index:
                    # 加入换手率
                    tr = stock_data.loc[self.trading_dates[day], 'turn']
                    sum_TR += tr
                else:
                    valid_days -= 1
            if valid_days > 2:
                aver_TR = sum_TR / valid_days
                tr1 = stock_data.loc[self.trading_dates[days[-1]], 'turn']
                tr2 = stock_data.loc[self.trading_dates[days[-1] - 1], 'turn']
                ratio = (tr1 + tr2) / (aver_TR * 2)
            else:
                ratio = 0
            if ratio > 0:
                stocks_data.loc[stock_code, 'ratio'] = ratio

        # print(stocks_data)
        stocks_data.sort_values(by='ratio', ascending=False, inplace=True)
        # print(stocks_data.index)
        if len(stocks_data.index) > number:
            # 取0到number-1共number只股票
            return stocks_data.index[0:number]
        else:
            # 取全部股票
            return stocks_data.index[:]

    def __nn_choose(self, model_type: str, today: tuple, number: int):
        """
        选择指定NN预测未来data_days天涨幅最高的number只股票
        """
        # 第一次买入策略应大于策略所需数据天数
        if today[1] < self.data_days:
            return pd.Series(None)
        # 建立用于计算预测涨跌幅的DF
        stocks_data = pd.DataFrame(self.stocks_codes)
        stocks_data['change'] = 0
        stocks_data = stocks_data.set_index('code')
        avail_num = 0
        # 预测每只股票未来data_days天涨跌幅
        for stock_code in self.stocks_codes:
            change = getattr(self.prediction, 'predict_' + model_type)(stock_code, today)
            if type(change) != int:
                # tensor直接取值
                change = change[0, 0].item()
            if change > 0:
                # 去除小于0的预测值
                stocks_data.loc[stock_code, 'change'] = change
                avail_num += 1
        # 排序
        stocks_data.sort_values(by='change', ascending=False, inplace=True)
        if avail_num > number:
            # 取0到number-1共number只股票
            return stocks_data.index[0:number]
        else:
            # 取全部有效股票
            return stocks_data.index[:avail_num]

    def choose_by_cnn(self, today: tuple, number: int):
        return self.__nn_choose('cnn', today, number)

    def choose_by_lstm(self, today: tuple, number: int):
        return self.__nn_choose('lstm', today, number)

    def choose_by_gru(self, today: tuple, number: int):
        return self.__nn_choose('gru', today, number)

    def choose_by_rnn_tanh(self, today: tuple, number: int):
        return self.__nn_choose('rnn_tanh', today, number)

    def choose_by_rnn_relu(self, today: tuple, number: int):
        return self.__nn_choose('rnn_relu', today, number)

    def choose_by_resnet18(self, today: tuple, number: int):
        return self.__nn_choose('resnet18', today, number)

    def choose_by_resnet34(self, today: tuple, number: int):
        return self.__nn_choose('resnet34', today, number)

    def choose_by_resnet50(self, today: tuple, number: int):
        return self.__nn_choose('resnet50', today, number)

    def choose_by_resnet101(self, today: tuple, number: int):
        return self.__nn_choose('resnet101', today, number)

    def choose_by_resnet152(self, today: tuple, number: int):
        return self.__nn_choose('resnet152', today, number)

    def choose_by_densenet(self, today: tuple, number: int):
        return self.__nn_choose('densenet', today, number)

    def choose_by_ensemble(self, today: tuple):
        chosen_num = pd.DataFrame()
        chosen_num['code'] = self.stocks_codes
        chosen_num['num'] = 0
        chosen_num.set_index('code', inplace=True)
        number = 300
        for chosen in (self.__nn_choose('cnn', today, number),
                       self.__nn_choose('lstm', today, number),
                       self.__nn_choose('rnn_relu', today, number),
                       self.__nn_choose('resnet18', today, number),
                       self.__nn_choose('resnet34', today, number),
                       self.__nn_choose('resnet50', today, number),
                       self.__nn_choose('densenet', today, number)):
            for stock_code in self.stocks_codes:
                if stock_code in chosen:
                    chosen_num.loc[stock_code, 'num'] += 1
        return chosen_num[chosen_num['num'] >= 3].index
